package com.kapukapu.authorization.config;

import com.kapukapu.authorization.util.UserInfoHodler;
import lombok.SneakyThrows;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.config.annotation.authentication.configuration.AuthenticationConfiguration;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.config.annotation.configurers.ClientDetailsServiceConfigurer;
import org.springframework.security.oauth2.config.annotation.web.configuration.AuthorizationServerConfigurerAdapter;
import org.springframework.security.oauth2.config.annotation.web.configuration.AuthorizationServerEndpointsConfiguration;
import org.springframework.security.oauth2.config.annotation.web.configurers.AuthorizationServerEndpointsConfigurer;
import org.springframework.security.oauth2.config.annotation.web.configurers.AuthorizationServerSecurityConfigurer;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenEnhancerChain;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.oauth2.provider.token.store.JdbcTokenStore;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;

import javax.sql.DataSource;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

/**
 * @author: sivan
 * @date: 2020/03/03
 * @email: zifanuu@gmail.com
 */
@Configuration
@Import(AuthorizationServerEndpointsConfiguration.class)
public class JwtOAuth2AuthorizationConfig extends AuthorizationServerConfigurerAdapter {
    @Value("${security.oauth2.private-key}")
    private String privateKey;
    @Value("${security.oauth2.public-key}")
    private String publicKey;

    private AuthenticationManager authenticationManager;
    private PasswordEncoder passwordEncoder;
    private DataSource dataSource;

    public JwtOAuth2AuthorizationConfig(
        AuthenticationConfiguration authenticationConfiguration,
        PasswordEncoder passwordEncoder, DataSource dataSource) throws Exception {
        this.authenticationManager = authenticationConfiguration
            .getAuthenticationManager();
        this.passwordEncoder = passwordEncoder;
        this.dataSource = dataSource;
    }

    @Override
    public void configure(AuthorizationServerSecurityConfigurer security)
        throws Exception {
        super.configure(security);
        security.tokenKeyAccess("permitAll()").checkTokenAccess("permitAll()")
            .allowFormAuthenticationForClients();
    }

    @Override
    public void configure(ClientDetailsServiceConfigurer clients) throws Exception {
        clients.jdbc(this.dataSource).passwordEncoder(this.passwordEncoder);
    }

    @Override
    public void configure(AuthorizationServerEndpointsConfigurer endpoints) {
        endpoints.authenticationManager(this.authenticationManager)
            .accessTokenConverter(accessTokenConverter())
            .tokenEnhancer(tokenEnhancerChain())
            .tokenStore(tokenStore());
    }

    private TokenEnhancer tokenEnhancerChain() {
        TokenEnhancerChain tokenEnhancerChain = new TokenEnhancerChain();
        tokenEnhancerChain.setTokenEnhancers(Arrays.asList(
            JwtTokenEnhancerFactory.userInfoTokenEnhancer(), accessTokenConverter()));
        return tokenEnhancerChain;
    }

    @Bean
    public TokenStore tokenStore() {
        return new JdbcTokenStore(dataSource);
    }

    @Bean
    @SneakyThrows
    public JwtAccessTokenConverter accessTokenConverter() {
        JwtAccessTokenConverter converter = new JwtAccessTokenConverter();

        converter.setSigningKey(privateKey);
        converter.setVerifierKey(publicKey);
        return converter;
    }

    static class JwtTokenEnhancerFactory {

        protected static TokenEnhancer userInfoTokenEnhancer() {
            return (accessToken, authentication) -> {
                final Map<String, Object> additionalInfo = new HashMap<>(1);
                additionalInfo.put("name", UserInfoHodler.getUserInfo().getName());
                additionalInfo.put("user_id", UserInfoHodler.getUserInfo().getId());

                UserInfoHodler.destroy();

                ((DefaultOAuth2AccessToken) accessToken)
                    .setAdditionalInformation(additionalInfo);
                return accessToken;
            };
        }

    }
}
